package com.lc1993929.mr.flowlog.enhance;

import org.apache.hadoop.fs.FSDataOutputStream;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.io.NullWritable;
import org.apache.hadoop.io.Text;
import org.apache.hadoop.mapreduce.RecordWriter;
import org.apache.hadoop.mapreduce.TaskAttemptContext;
import org.apache.hadoop.mapreduce.lib.output.FileOutputFormat;

import java.io.IOException;

/**
 * Created by LiuChang on 2018/1/8.
 */
public class EnhanceOutputFormat extends FileOutputFormat<Text, NullWritable> {

    @Override
    public RecordWriter<Text, NullWritable> getRecordWriter(TaskAttemptContext taskAttemptContext) throws IOException, InterruptedException {
        FileSystem fileSystem = FileSystem.get(taskAttemptContext.getConfiguration());
        FSDataOutputStream tocrawlOs = fileSystem.create(new Path("hdfs://hdp1:9000/tocrawl/result.dat"));
        FSDataOutputStream enhanceOs = fileSystem.create(new Path("hdfs://hdp1:9000/enhance/result.dat"));
        return new EnhanceRecordWriter(tocrawlOs, enhanceOs);
    }


    static class EnhanceRecordWriter extends RecordWriter<Text, NullWritable> {

        FSDataOutputStream tocrawlOs = null;
        FSDataOutputStream enhanceOs = null;

        public EnhanceRecordWriter(FSDataOutputStream tocrawlOs, FSDataOutputStream enhanceOs) {
            this.tocrawlOs = tocrawlOs;
            this.enhanceOs = enhanceOs;
        }

        @Override
        public void write(Text text, NullWritable nullWritable) throws IOException, InterruptedException {

            if (text.toString().contains("\001tocrawl")) {
                //    写入待爬清单
                tocrawlOs.write(text.toString().getBytes());
                tocrawlOs.write("\n".getBytes());
                tocrawlOs.flush();
            } else {
                //    写入增强结果日志文件
                enhanceOs.write(text.toString().getBytes());
                enhanceOs.write("\n".getBytes());
                enhanceOs.flush();
            }
        }

        @Override
        public void close(TaskAttemptContext taskAttemptContext) throws IOException, InterruptedException {
            if (tocrawlOs != null) {
                tocrawlOs.close();
            }
            if (enhanceOs != null) {
                enhanceOs.close();
            }
        }
    }
}
